package com.wfg.flink.connector.mongodb;


import com.alibaba.fastjson2.JSON;
import com.mongodb.client.model.InsertOneModel;
import com.wfg.flink.connector.mongodb.model.WellCastingInfo;
import com.wfg.flink.connector.mongodb.schema.WellCastingInfoDeserializationSchema;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.connector.mongodb.sink.MongoSink;
import org.apache.flink.connector.mongodb.source.MongoSource;
import org.apache.flink.connector.mongodb.source.enumerator.splitter.PartitionStrategy;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.bson.BsonDocument;


/**
 * @author wfg
 */
@Slf4j
public class Main {
    public static void main(String[] args) throws Exception {


        MongoSource<WellCastingInfo> mongoSource = MongoSource.<WellCastingInfo>builder()
                .setUri("mongodb://root:123456@127.0.0.1:27017,127.0.0.1:27018,127.0.0.1:27019/admin?replicaSet=rs0&authSource=admin")
                .setDatabase("sjzz")
                .setCollection("wellCastingInfo")
//                .setProjectedFields("_id", "f0", "f1")
                .setFetchSize(2048)
                .setLimit(10000)
                .setNoCursorTimeout(true)
                .setPartitionStrategy(PartitionStrategy.SAMPLE)
                .setPartitionSize(MemorySize.ofMebiBytes(64))
                .setSamplesPerPartition(10)
                .setDeserializationSchema(new WellCastingInfoDeserializationSchema())
                .build();

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // 从 MongoDB 读取数据
        DataStream<WellCastingInfo> sourceStream = env.fromSource(mongoSource, WatermarkStrategy.noWatermarks(), "Mongo Source");

        // 进行转换（如果需要）
        DataStream<WellCastingInfo> transformedStream = sourceStream.map((MapFunction<WellCastingInfo, WellCastingInfo>) value -> {
            // 转换逻辑
            return value;
        });

        MongoSink<WellCastingInfo> sink = MongoSink.<WellCastingInfo>builder()
                .setUri("mongodb://root:123456@127.0.0.1:27017,127.0.0.1:27018,127.0.0.1:27019/admin?replicaSet=rs0&authSource=admin")
                .setDatabase("sjzz")
                .setCollection("wellCastingInfo_back")
                .setMaxRetries(3)
//                .setDeliveryGuarantee(DeliveryGuarantee.AT_LEAST_ONCE)
                .setSerializationSchema(
                        (input, context) -> new InsertOneModel<>(BsonDocument.parse(JSON.toJSONString(input))))
                .build();
        transformedStream.sinkTo(sink);
//        stream.sinkTo(sink);

        // 执行作业
        env.execute("Mongo Flink Demo");

    }

}
